package com.lfg.learn.rbtree;

import com.alibaba.fastjson.JSON;

import java.util.ArrayList;
import java.util.List;

public class RBTree {

    Node root = null;

    public static void main(String[] args) {
        RBTree tree = new RBTree();
        tree.insert(1);
        System.out.println(JSON.toJSONString(tree.visit(tree.root)));
        tree.insert(2);
        System.out.println(JSON.toJSONString(tree.visit(tree.root)));
        tree.insert(3);
        System.out.println(JSON.toJSONString(tree.visit(tree.root)));
        tree.insert(4);
        System.out.println(JSON.toJSONString(tree.visit(tree.root)));
        tree.insert(5);
        System.out.println(JSON.toJSONString(tree.visit(tree.root)));
        tree.insert(6);
        System.out.println(JSON.toJSONString(tree.visit(tree.root)));
        tree.insert(7);
        System.out.println(JSON.toJSONString(tree.visit(tree.root)));
        tree.insert(8);
        System.out.println(JSON.toJSONString(tree.visit(tree.root)));


        tree.delete(4);
        tree.valid(tree.root,0);
        System.out.println(JSON.toJSONString(tree.visit(tree.root)));

        tree.delete(7);
        tree.valid(tree.root,0);
        System.out.println(JSON.toJSONString(tree.visit(tree.root)));

        tree.delete(3);
        tree.valid(tree.root,0);
        System.out.println(JSON.toJSONString(tree.visit(tree.root)));

        tree.delete(2);
        tree.valid(tree.root,0);
        System.out.println(JSON.toJSONString(tree.visit(tree.root)));
    }


    public List<String> visit(Node val) {
        List<String> ret = new ArrayList<>();
        ret.add(val.value + " " + val.black + ", ");
        if (null != val.left) {
            ret.addAll(visit(val.left));
        }
        if (null != val.right) {
            ret.addAll(visit(val.right));
        }
        return ret;
    }

    public void delete(int val) {
        Node node = search(val);
        if (node == null) {
            return;
        }

        Node successor = successor(node);
        swapVal(node, successor);
        balanceAfterDelete(successor);
        deleteNode(successor);

    }

    public void deleteNode(Node node) {
        Node p = node.parent;
        Node l = node.left;
        Node r = node.right;
        Node replace = l == null ? r : l;

        if (p == null) {
            root = replace;
        } else {
            if (p.left == node) {
                p.left = replace;
            } else {
                p.right = replace;
            }
        }


    }

    private void balanceAfterDelete(Node x) {

        /**
         * 1.有子节点，则从子节点替换上来
         * 2.兄弟节点是红色，则旋转后父节点改为红色。
         * 3.兄弟节点是黑色，则看兄弟的2个子节点都是黑色，则将兄弟节点改为红色，以父节点为平衡点
         * 4.兄弟节点是黑色，则看兄弟的远节点是红色，旋转后改为父节点和兄弟的远子节点都改为黑色， 兄弟节点改成父节点的颜色
         * 5.兄弟节点是黑色，则看兄弟的远节点是黑色，先对兄弟节点旋转，然后以父节点旋转后改为父节点和兄弟的远子节点都改为黑色， 兄弟节点改成父节点的颜色
         */
        if (!x.black) {
            return;
        }
        while (x != root) {
            Node p = x.parent;
            Node pp = p.parent;
            if (!x.black) {
                break;
            }
            if (x.right != null) {
                x.right.black = true;
                x = root;
                continue;
            }

            if (x == p.left) {
                Node pr = p.right;
                if (!pr.black) {
                    leftRound(p);
                    pr.black = true;
                    x = root;
                } else {
                    Node prr = pr.right;
                    Node prl = pr.left;
                    if ((prr == null || prr.black) && (prl == null || prl.black)) {
                        pr.black = false;
                        x = p;
                    } else {
                        if (prr == null || prr.black) {
                            rightRound(pr);
                            pr.black = false;
                            prl.black = true;
                            pr = p.right;
                            prr = pr.right;
                        }
                        leftRound(p);
                        pr.black = p.black;
                        prr.black = true;
                        p.black = true;
                        x = root;
                    }
                }
            } else {
                Node pl = p.left;
                if (!pl.black) {
                    rightRound(p);
                    pl.black = true;
                    x = root;
                } else {
                    Node plr = pl.right;
                    Node pll = pl.left;
                    if ((plr == null || plr.black) && (pll == null || pll.black)) {
                        pl.black = false;
                        x = p;
                    } else {
                        if (plr == null || plr.black) {
                            leftRound(pl);
                            pl.black = false;
                            pll.black = true;
                            pl = p.right;
                            pll = pl.left;
                        }
                        rightRound(p);
                        pl.black = p.black;
                        p.black = true;
                        pll.black = true;
                        x = root;
                    }
                }
            }
        }
        x.black = true;

    }

    private void swapVal(Node node, Node successor) {
        if (node == successor) {
            return;
        }
        int val = node.value;

        node.value = successor.value;
        successor.value = val;
    }

    public Node successor(Node node) {
        if (node == null) {
            return null;
        }
        Node ret = node.right;
        if (null == ret) {
            return node;
        }
        while (ret.left != null) {
            ret = ret.left;
        }
        return ret;
    }

    public Node search(int val) {
        Node node = root;
        if (node == null) {
            return null;
        }
        while (node != null) {
            if (node.value == val) {
                return node;
            } else if (val < node.value) {
                node = node.left;
            } else {
                node = node.right;
            }
        }
        return node;
    }


    public void insert(int val) {
        Node node = root;
        Node newNode = new Node();
        newNode.value = val;
        if (node == null) {
            root = newNode;
            root.black = true;
            return;
        }
        while (node != null) {
            if (node.value == val) {
                break;
            } else if (val < node.value) {
                if (null == node.left) {
                    node.left = newNode;
                    newNode.parent = node;
                    break;
                } else {
                    node = node.left;
                }
            } else {
                if (null == node.right) {
                    node.right = newNode;
                    newNode.parent = node;
                    break;
                } else {
                    node = node.right;
                }
            }
        }
        balanceAfterInsert(newNode);
    }


    public void balanceAfterInsert(Node node) {
        // 父节点是黑色，则挑过
        // 父节点是红色，兄弟节点是黑色，则移动到兄弟节点
        // 兄弟节点和父节点都是红色， 则兄弟节点和当前节点设置为黑色，移动到父节点平衡

        while (node.parent != null) {
            Node p = node.parent;
            if (p.black) {
                return;
            }
            Node pp = node.parent.parent;
            if (pp == null) {
                p.black = true;
                return;
            }
            if (pp.left == node.parent) {
                Node right = pp.right;
                if (right == null || right.black) {
                    rightRound(pp);
                    p.black = true;
                    pp.black = false;
                } else {
                    p.black = true;
                    right.black = true;
                    pp.black = false;
                    node = pp;
                }
            } else {
                Node pl = pp.left;
                if (pl == null || pl.black) {
                    leftRound(pp);
                    p.black = true;
                    pp.black = false;
                } else {
                    p.black = true;
                    pl.black = true;
                    pp.black = false;
                    node = pp;
                }
            }
        }
        node.black = true;
    }

    public void rightRound(Node node) {
        Node pp = node.parent;
        Node p = node.left;
        Node right = p.right;

        p.right = node;
        node.parent = p;

        node.left = right;
        if(null != right) {
            right.parent = node;
        }
        if (pp != null) {
            if (pp.left == node) {
                pp.left = p;
            } else {
                pp.right = p;
            }
        } else {
            root = p;
        }
        p.parent = pp;

    }


    public void leftRound(Node node) {
        Node pp = node.parent;

        Node p = node.right;
        Node left = p.left;

        p.left = node;
        node.parent = p;

        node.right = left;
        if (left != null) {
            left.parent = node;
        }
        if (pp != null) {
            if (pp.left == node) {
                pp.left = p;
            } else {
                pp.right = p;
            }
        } else {
            root = p;
        }
        p.parent = pp;
    }

    public void valid(Node node, int deep){
        if(node == null){
            System.out.println("deep=" + deep);
            return;
        }
        if(node.black){
            deep++;
        }
        valid(node.left, deep);
        valid(node.right, deep);

    }

}
